import numpy as np
import random
from typing import List
import gym
import gym.spaces as spaces
from typing import Dict, Tuple, Union
from collections import OrderedDict
import numba


class NetworkEnv(gym.core.Env):

    def __init__(self,
                 num_transmitter: int = 3,
                 num_receiver: int = 3,
                 num_channel: int = 2,
                 bmax: int = 10,
                 emax: int = 10,
                 link_probabilities: Union[List[float], np.ndarray] = [
                     0.1, 0.2, 0.5, 0.2],
                 eng_probabilities: Union[List[float], np.ndarray] = [
                     0.204, 0.338, 0.266, 0.192],
                 max_episode_steps: int = 50):
        """ 创建网络通讯gym环境

        :param num_Transmitter: 发送端用户数
        :type num_Transmitter: int
        :param num_Receiver: 接收端用户数
        :type num_Receiver: int
        :param num_Channel: 信道数
        :type num_Channel: int
        :param bmax: 每个块储存数据
        :type bmax: int
        :param emax: R储存能量的最大值
        :type emax: int
        :param link_probabilities: 链路状态分布
        :type link_probabilities: Union[List[float], np.ndarray]
        :param eng_probabilities: 能量吸收模型
        :type eng_probabilities: Union[List[float], np.ndarray]
        :param max_episode_steps: 每回合最大步数
        :type max_episode_steps: int
        """
        assert np.isclose(sum(link_probabilities), 1.) and\
            len(link_probabilities) == 4
        assert num_transmitter > 0
        assert num_receiver > 0
        assert num_channel > 0
        assert max_episode_steps > 0
        self.num_transmitter = num_transmitter  # 发送端用户数
        self.num_receiver = num_receiver  # 接收端用户数
        self.num_channel = num_channel  # 信道数
        self.bmax = bmax  # 每个块存储数据
        self.emax = emax  # R存储能量的最大值
        self.link_probabilities = np.asarray(link_probabilities)  # 链路状态分布
        self.link_probabilities_cdf = NetworkEnv._get_cdf(
            self.link_probabilities)
        self.eng_probabilities = eng_probabilities  # 能量吸收模型

        self.rest_data_b1 = 0  # 第一对用户的数据缓存池中的剩余数据量
        self.rest_data_b2 = 0  # 第二对用户的数据缓存池中的剩余数据量
        self.rest_data_b3 = 0  # 第三对用户的数据缓存池中的剩余数据量
        self.rest_power = 0  # 整个系统当前时刻的剩余能量
        self.rw_u = 0.5  # 上传链路的reward因子
        self.rw_d = 2.5  # 下载链路的reward因子
        self.reward = 0
        self.harvest_power = 0  # 当前时刻从外界可以吸收的能量
        self.consume_power = 1  # 选择转发的动作时消耗的能量（固定）

        # 环境状态动作描述
        self.action_space = spaces.Discrete(25)
        self.observation_space = spaces.Dict({
            'observation': spaces.Box(low=0., high=1., shape=(
                2*self.num_transmitter*self.num_channel+4,)),
            'link_state': spaces.Box(
                low=0, high=3,
                shape=(2*self.num_transmitter*self.num_channel,)),
            'rest_datab1': spaces.Box(low=0, high=3, shape=(1,)),
            'rest_datab2': spaces.Box(low=0, high=3, shape=(1,)),
            'rest_datab3': spaces.Box(low=0, high=3, shape=(1,)),
            'rest_power': spaces.Box(low=0, high=3, shape=(1,)),
        })
        self.max_step = max_episode_steps
        self.step_counter = 0
        self.done = None

        # utils
        self._value_list = [0, 1, 2, 3]

    @staticmethod
    def _get_cdf(probs):
        """计算累计分布
        """
        tri_low = np.tril(np.ones([len(probs), len(probs)]))  # 下三角矩阵
        cdf = tri_low.dot(probs)
        return cdf

    def change_link_prob(self, prob: Union[List[float], np.ndarray]):
        assert np.isclose(sum(prob), 1.) and len(prob) == 4
        self.link_probabilities = np.array(prob)
        self.link_probabilities_cdf = NetworkEnv._get_cdf(
            self.link_probabilities)

    def change_eng_prob(self, prob: Union[List[float], np.ndarray]):
        assert np.isclose(sum(prob), 1.) and len(prob) == 4
        self.eng_probabilities = np.asarray(prob)

    def _change_harvest_power(self):
        # 环境每次交互后都会改变其能量状态
        self.harvest_power = np.random.choice(
            self._value_list, p=self.eng_probabilities)

    def reset(self) -> Tuple[dict, dict]:
        link_state = np.random.randint(
            0, 4, size=2 * self.num_transmitter * self.num_channel)
        rest_state = np.random.randint(1, 4, size=4)

        # update state
        self.link_state = link_state
        self.rest_data_b1 = rest_state[0]
        self.rest_data_b2 = rest_state[1]
        self.rest_data_b3 = rest_state[2]
        self.rest_power = rest_state[3]
        obs = np.concatenate([link_state/3, rest_state/10], axis=0)
        self._change_harvest_power()
        self.step_counter = 0
        self.done = False

        obs = dict(
            observation=obs,
            link_state=link_state,
            rest_datab1=rest_state[0],
            rest_datab2=rest_state[1],
            rest_datab3=rest_state[2],
            rest_power=rest_state[3]
        )
        return obs

    def render(self, mode='human'):
        raise NotImplementedError()

    def close(self):
        pass

    def step(self, action: int):
        if self.done:
            raise RuntimeError(
                "The environment has done, expect to run reset()")

        # UE1->uplink1, UE2->uplink2
        if action == 0:
            # self.transmit_data_amount = 0
            # 选择通过uplink1接收UE1的数据包后b1的数据包数量的变化
            self.new_rest_data_b1 = min(
                self.rest_data_b1 + min(self.link_state[0],
                                        self.bmax - self.rest_data_b1),
                self.bmax)
            # 选择通过uplink2接收UE2的数据包后b2的数据包数量的变化
            self.new_rest_data_b2 = min(
                self.rest_data_b2 + min(self.link_state[3],
                                        self.bmax - self. rest_data_b2),
                self.bmax)
            self.actual_harvest_power = min(
                self.harvest_power, self.emax - self.rest_power)  # 实际吸收的能量
            self.rest_power = min(
                self.rest_power + self.actual_harvest_power, self.emax)
            self.reward = self.rw_u *\
                (self.new_rest_data_b1 - self.rest_data_b1) + \
                self.rw_u * (self.new_rest_data_b2 - self.rest_data_b2)
            self.rest_data_b1 = self.new_rest_data_b1
            self.rest_data_b2 = self.new_rest_data_b2

        # UE1->uplink2, UE2->uplink1
        if action == 1:
            # self.transmit_data_amount = 0
            # 选择通过uplink2接收UE1的数据包后b1的数据包数量的变化
            self.new_rest_data_b1 = min(
                self.rest_data_b1 + min(self.link_state[1],
                                        self.bmax - self.rest_data_b1),
                self.bmax)
            # 选择通过uplink1接收UE2的数据包后b2的数据包数量的变化
            self.new_rest_data_b2 = min(
                self.rest_data_b2 + min(
                    self.link_state[2], self.bmax - self.rest_data_b2),
                self.bmax)
            self.actual_harvest_power = min(
                self.harvest_power, self.emax - self.rest_power)  # 实际吸收的能量
            self.rest_power = min(
                self.rest_power + self.actual_harvest_power, self.emax)
            self.reward = self.rw_u * \
                (self.new_rest_data_b1 - self.rest_data_b1) + \
                self.rw_u * (self.new_rest_data_b2 - self.rest_data_b2)
            self.rest_data_b1 = self.new_rest_data_b1
            self.rest_data_b2 = self.new_rest_data_b2

        # UE1->uplink1, UE3->uplink2
        if action == 2:
            # self.transmit_data_amount = 0
            # 选择通过uplink1接收UE1的数据包后b1的数据包数量的变化
            self.new_rest_data_b1 = min(
                self.rest_data_b1 + min(
                    self.link_state[0], self.bmax - self.rest_data_b1),
                self.bmax)
            # 选择通过uplink2接收UE3的数据包后b3的数据包数量的变化
            self.new_rest_data_b3 = min(
                self.rest_data_b3 + min(
                    self.link_state[5], self.bmax - self.rest_data_b3),
                self.bmax)
            self.actual_harvest_power = min(
                self.harvest_power, self.emax - self.rest_power)  # 实际吸收的能量
            self.rest_power = min(
                self.rest_power + self.actual_harvest_power, self.emax)
            self.reward = self.rw_u *\
                (self.new_rest_data_b1 - self.rest_data_b1) + \
                self.rw_u * (self.new_rest_data_b3 - self.rest_data_b3)
            self.rest_data_b1 = self.new_rest_data_b1
            self.rest_data_b3 = self.new_rest_data_b3

        # UE1->uplink2, UE3->uplink1
        if action == 3:
            # self.transmit_data_amount = 0
            # 选择通过uplink2接收UE1的数据包后b1的数据包数量的变化
            self.new_rest_data_b1 = min(
                self.rest_data_b1 + min(
                    self.link_state[1], self.bmax - self.rest_data_b1),
                self.bmax)
            # 选择通过uplink1接收UE3的数据包后b3的数据包数量的变化
            self.new_rest_data_b3 = min(
                self.rest_data_b3 + min(
                    self.link_state[4], self.bmax - self.rest_data_b3),
                self.bmax)
            self.actual_harvest_power = min(
                self.harvest_power, self.emax - self.rest_power)  # 实际吸收的能量
            self.rest_power = min(
                self.rest_power + self.actual_harvest_power, self.emax)
            self.reward = self.rw_u * \
                (self.new_rest_data_b1 - self.rest_data_b1) + \
                self.rw_u * (self.new_rest_data_b3 - self.rest_data_b3)
            self.rest_data_b1 = self.new_rest_data_b1
            self.rest_data_b3 = self.new_rest_data_b3

        # UE2->uplink1, UE3->uplink2
        if action == 4:
            # self.transmit_data_amount = 0
            # 选择通过uplink1接收UE2的数据包后b2的数据包数量的变化
            self.new_rest_data_b2 = min(
                self.rest_data_b2 + min(
                    self.link_state[2], self.bmax - self.rest_data_b2),
                self.bmax)
            # 选择通过uplink2接收UE3的数据包后b3的数据包数量的变化
            self.new_rest_data_b3 = min(
                self.rest_data_b3 + min(
                    self.link_state[5], self.bmax - self.rest_data_b3),
                self.bmax)
            self.actual_harvest_power = min(
                self.harvest_power, self.emax - self.rest_power)  # 实际吸收的能量
            self.rest_power = min(
                self.rest_power + self.actual_harvest_power, self.emax)
            self.reward = self.rw_u *\
                (self.new_rest_data_b2 - self.rest_data_b2) + \
                self.rw_u * (self.new_rest_data_b3 - self.rest_data_b3)
            self.rest_data_b2 = self.new_rest_data_b2
            self.rest_data_b3 = self.new_rest_data_b3

        # UE2->uplink2, UE3->uplink1
        if action == 5:
            # self.transmit_data_amount = 0
            # 选择通过uplink2接收UE2的数据包后b2的数据包数量的变化
            self.new_rest_data_b2 = min(
                self.rest_data_b2 + min(self.link_state[3], self.bmax - self.rest_data_b2), self.bmax)
            # 选择通过uplink1接收UE3的数据包后b3的数据包数量的变化
            self.new_rest_data_b3 = min(
                self.rest_data_b3 + min(self.link_state[4], self.bmax - self.rest_data_b3), self.bmax)
            self.actual_harvest_power = min(
                self.harvest_power, self.emax - self.rest_power)  # 实际吸收的能量
            self.rest_power = min(
                self.rest_power + self.actual_harvest_power, self.emax)
            self.reward = self.rw_u * (self.new_rest_data_b2 - self.rest_data_b2) + \
                self.rw_u * (self.new_rest_data_b3 - self.rest_data_b3)
            self.rest_data_b2 = self.new_rest_data_b2
            self.rest_data_b3 = self.new_rest_data_b3

        # DE1->downlink1, DE2->downlink2
        if action == 6 and self.rest_power >= 2:
            # self.transmit_data_amount = min(self.link_state[2] + self.link_state[3], self.rest_power)  # RN能够转发给DE2的数据包的数量
            # 选择通过downlink1转发给DE1数据包后b1的数据包数量的变化
            self.new_rest_data_b1 = max(
                self.rest_data_b1 - min(self.link_state[6], self.rest_data_b1, self.rest_power), 0)
            self.rest_power = self.rest_power - self.consume_power
            # 选择通过downlink2转发给DE2数据包后b2的数据包数量的变化
            self.new_rest_data_b2 = max(
                self.rest_data_b2 - min(self.link_state[9], self.rest_data_b2, self.rest_power), 0)
            self.rest_power = self.rest_power - self.consume_power
            self.actual_harvest_power = min(
                self.harvest_power, self.emax - self.rest_power)  # 实际吸收的能量
            self.rest_power = min(
                self.rest_power + self.actual_harvest_power, self.emax)
            self.reward = self.rw_d * (self.rest_data_b1 - self.new_rest_data_b1) + \
                self.rw_d * (self.rest_data_b2 - self.new_rest_data_b2)
            self.rest_data_b1 = self.new_rest_data_b1
            self.rest_data_b2 = self.new_rest_data_b2

        # DE1->downlink1, DE3->downlink2
        if action == 7 and self.rest_power >= 2:
            # self.transmit_data_amount = min(self.link_state[2] + self.link_state[3], self.rest_power)  # RN能够转发给DE2的数据包的数量
            # 选择通过downlink1转发给DE1数据包后b1的数据包数量的变化
            self.new_rest_data_b1 = max(
                self.rest_data_b1 - min(self.link_state[6], self.rest_data_b1, self.rest_power), 0)
            self.rest_power = self.rest_power - self.consume_power
            # 选择通过downlink2转发给DE3数据包后b3的数据包数量的变化
            self.new_rest_data_b3 = max(
                self.rest_data_b3 - min(self.link_state[11], self.rest_data_b3, self.rest_power), 0)
            self.rest_power = self.rest_power - self.consume_power
            self.actual_harvest_power = min(
                self.harvest_power, self.emax - self.rest_power)  # 实际吸收的能量
            self.rest_power = min(
                self.rest_power + self.actual_harvest_power, self.emax)
            self.reward = self.rw_d * (self.rest_data_b1 - self.new_rest_data_b1) + \
                self.rw_d * (self.rest_data_b3 - self.new_rest_data_b3)
            self.rest_data_b1 = self.new_rest_data_b1
            self.rest_data_b3 = self.new_rest_data_b3

        # DE1->downlink2, DE2->downlink1
        if action == 8 and self.rest_power >= 2:
            # self.transmit_data_amount = min(self.link_state[2] + self.link_state[3], self.rest_power)  # RN能够转发给DE2的数据包的数量
            # 选择通过downlink2转发给DE1数据包后b1的数据包数量的变化
            self.new_rest_data_b1 = max(
                self.rest_data_b1 - min(self.link_state[7], self.rest_data_b1, self.rest_power), 0)
            self.rest_power = self.rest_power - self.consume_power
            # 选择通过downlink1转发给DE2数据包后b2的数据包数量的变化
            self.new_rest_data_b2 = max(
                self.rest_data_b2 - min(self.link_state[8], self.rest_data_b2, self.rest_power), 0)
            self.rest_power = self.rest_power - self.consume_power
            self.actual_harvest_power = min(
                self.harvest_power, self.emax - self.rest_power)  # 实际吸收的能量
            self.rest_power = min(
                self.rest_power + self.actual_harvest_power, self.emax)
            self.reward = self.rw_d * (self.rest_data_b1 - self.new_rest_data_b1) + \
                self.rw_d * (self.rest_data_b2 - self.new_rest_data_b2)
            self.rest_data_b1 = self.new_rest_data_b1
            self.rest_data_b2 = self.new_rest_data_b2

        # DE1->downlink2, DE3->downlink1
        if action == 9 and self.rest_power >= 2:
            # self.transmit_data_amount = min(self.link_state[2] + self.link_state[3], self.rest_power)  # RN能够转发给DE2的数据包的数量
            # 选择通过downlink2转发给DE1数据包后b1的数据包数量的变化
            self.new_rest_data_b1 = max(
                self.rest_data_b1 - min(self.link_state[7], self.rest_data_b1, self.rest_power), 0)
            self.rest_power = self.rest_power - self.consume_power
            # 选择通过downlink1转发给DE3数据包后b3的数据包数量的变化
            self.new_rest_data_b3 = max(
                self.rest_data_b3 - min(self.link_state[10], self.rest_data_b3, self.rest_power), 0)
            self.rest_power = self.rest_power - self.consume_power
            self.actual_harvest_power = min(
                self.harvest_power, self.emax - self.rest_power)  # 实际吸收的能量
            self.rest_power = min(
                self.rest_power + self.actual_harvest_power, self.emax)
            self.reward = self.rw_d * (self.rest_data_b1 - self.new_rest_data_b1) + \
                self.rw_d * (self.rest_data_b3 - self.new_rest_data_b3)
            self.rest_data_b1 = self.new_rest_data_b1
            self.rest_data_b3 = self.new_rest_data_b3

        # DE2->downlink1, DE3->downlink2
        if action == 10 and self.rest_power >= 2:
            # self.transmit_data_amount = min(self.link_state[2] + self.link_state[3], self.rest_power)  # RN能够转发给DE2的数据包的数量
            # 选择通过downlink1转发给DE2数据包后b2的数据包数量的变化
            self.new_rest_data_b2 = max(
                self.rest_data_b2 - min(self.link_state[8], self.rest_data_b2, self.rest_power), 0)
            self.rest_power = self.rest_power - self.consume_power
            # 选择通过downlink2转发给DE3数据包后b3的数据包数量的变化
            self.new_rest_data_b3 = max(
                self.rest_data_b3 - min(self.link_state[11], self.rest_data_b3, self.rest_power), 0)
            self.rest_power = self.rest_power - self.consume_power
            self.actual_harvest_power = min(
                self.harvest_power, self.emax - self.rest_power)  # 实际吸收的能量
            self.rest_power = min(
                self.rest_power + self.actual_harvest_power, self.emax)
            self.reward = self.rw_d * (self.rest_data_b2 - self.new_rest_data_b2) + \
                self.rw_d * (self.rest_data_b3 - self.new_rest_data_b3)
            self.rest_data_b2 = self.new_rest_data_b2
            self.rest_data_b3 = self.new_rest_data_b3

        # DE2->downlink2, DE3->downlink1
        if action == 11 and self.rest_power >= 2:
            # self.transmit_data_amount = min(self.link_state[2] + self.link_state[3], self.rest_power)  # RN能够转发给DE2的数据包的数量
            self.new_rest_data_b2 = max(self.rest_data_b2 - min(
                self.link_state[9], self.rest_data_b2, self.rest_power), 0)  # 选择转发给UE2数据包后b2的数据包数量的变化
            self.rest_power = self.rest_power - self.consume_power
            self.new_rest_data_b3 = max(self.rest_data_b3 - min(
                self.link_state[10], self.rest_data_b3, self.rest_power), 0)  # 选择转发给UE2数据包后b2的数据包数量的变化
            self.rest_power = self.rest_power - self.consume_power
            self.actual_harvest_power = min(
                self.harvest_power, self.emax - self.rest_power)  # 实际吸收的能量
            self.rest_power = min(
                self.rest_power + self.actual_harvest_power, self.emax)
            self.reward = self.rw_d * (self.rest_data_b2 - self.new_rest_data_b2) + \
                self.rw_d * (self.rest_data_b3 - self.new_rest_data_b3)
            self.rest_data_b2 = self.new_rest_data_b2
            self.rest_data_b3 = self.new_rest_data_b3

        # RN choose to suspend the uplink1, UE1->uplink2
        if action == 12:
            # 选择通过uplink2接收UE1的数据包后b1的数据包数量的变化
            self.new_rest_data_b1 = min(
                self.rest_data_b1 + min(self.link_state[1], self.bmax - self.rest_data_b1), self.bmax)
            self.actual_harvest_power = min(
                self.harvest_power, self.emax - self.rest_power)  # 实际吸收的能量
            self.rest_power = min(
                self.rest_power + self.harvest_power, self.emax)
            self.reward = self.rw_u * \
                (self.new_rest_data_b1 - self.rest_data_b1)
            self.rest_data_b1 = self.new_rest_data_b1

        # RN choose to suspend the uplink1, UE2->uplink2
        if action == 13:
            # 选择通过uplink2接收UE3的数据包后b1的数据包数量的变化
            self.new_rest_data_b2 = min(
                self.rest_data_b2 + min(self.link_state[3], self.bmax - self.rest_data_b2), self.bmax)
            self.actual_harvest_power = min(
                self.harvest_power, self.emax - self.rest_power)  # 实际吸收的能量
            self.rest_power = min(
                self.rest_power + self.harvest_power, self.emax)
            self.reward = self.rw_u * \
                (self.new_rest_data_b2 - self.rest_data_b2)
            self.rest_data_b2 = self.new_rest_data_b2

        # RN choose to suspend the uplink1, UE3->uplink2
        if action == 14:
            # 选择通过uplink2接收UE3的数据包后b1的数据包数量的变化
            self.new_rest_data_b3 = min(
                self.rest_data_b3 + min(self.link_state[5], self.bmax - self.rest_data_b3), self.bmax)
            self.actual_harvest_power = min(
                self.harvest_power, self.emax - self.rest_power)  # 实际吸收的能量
            self.rest_power = min(
                self.rest_power + self.harvest_power, self.emax)
            self.reward = self.rw_u * \
                (self.new_rest_data_b3 - self.rest_data_b3)
            self.rest_data_b3 = self.new_rest_data_b3

        # RN choose to suspend the uplink2, UE1->uplink1
        if action == 15:
            # 选择通过uplink1接收UE1的数据包后b1的数据包数量的变化
            self.new_rest_data_b1 = min(
                self.rest_data_b1 + min(self.link_state[0], self.bmax - self.rest_data_b1), self.bmax)
            self.actual_harvest_power = min(
                self.harvest_power, self.emax - self.rest_power)  # 实际吸收的能量
            self.rest_power = min(
                self.rest_power + self.harvest_power, self.emax)
            self.reward = self.rw_u * \
                (self.new_rest_data_b1 - self.rest_data_b1)
            self.rest_data_b1 = self.new_rest_data_b1

        # RN choose to suspend the uplink2, UE2->uplink1
        if action == 16:
            # 选择通过uplink1接收UE2的数据包后b1的数据包数量的变化
            self.new_rest_data_b2 = min(
                self.rest_data_b2 + min(self.link_state[2], self.bmax - self.rest_data_b2), self.bmax)
            self.actual_harvest_power = min(
                self.harvest_power, self.emax - self.rest_power)  # 实际吸收的能量
            self.rest_power = min(
                self.rest_power + self.harvest_power, self.emax)
            self.reward = self.rw_u * \
                (self.new_rest_data_b2 - self.rest_data_b2)
            self.rest_data_b2 = self.new_rest_data_b2

        # RN choose to suspend the uplink2, UE3->uplink1
        if action == 17:
            # 选择通过uplink1接收UE3的数据包后b1的数据包数量的变化
            self.new_rest_data_b3 = min(
                self.rest_data_b3 + min(self.link_state[4], self.bmax - self.rest_data_b3), self.bmax)
            self.actual_harvest_power = min(
                self.harvest_power, self.emax - self.rest_power)  # 实际吸收的能量
            self.rest_power = min(
                self.rest_power + self.harvest_power, self.emax)
            self.reward = self.rw_u * \
                (self.new_rest_data_b3 - self.rest_data_b3)
            self.rest_data_b3 = self.new_rest_data_b3

        # RN choose to suspend the downlink1, DE1->downlink2
        if action == 18 and self.rest_power >= 1:
            # 选择通过downlink2转发给DE1数据包后b1的数据包数量的变化
            self.new_rest_data_b1 = max(
                self.rest_data_b1 - min(self.link_state[7], self.rest_data_b1, self.rest_power), 0)
            self.actual_harvest_power = min(
                self.harvest_power, self.emax - self.rest_power)  # 实际吸收的能量
            self.rest_power = min(
                self.rest_power + self.harvest_power - self.consume_power, self.emax)
            self.reward = self.rw_d * \
                (self.rest_data_b1 - self.new_rest_data_b1)
            self.rest_data_b1 = self.new_rest_data_b1

        # RN choose to suspend the downlink1, DE2->downlink2
        if action == 19 and self.rest_power >= 1:
            # 选择通过downlink2转发给DE2数据包后b2的数据包数量的变化
            self.new_rest_data_b2 = max(
                self.rest_data_b2 - min(self.link_state[9], self.rest_data_b2, self.rest_power), 0)
            self.actual_harvest_power = min(
                self.harvest_power, self.emax - self.rest_power)  # 实际吸收的能量
            self.rest_power = min(
                self.rest_power + self.harvest_power - self.consume_power, self.emax)
            self.reward = self.rw_d * \
                (self.rest_data_b2 - self.new_rest_data_b2)
            self.rest_data_b2 = self.new_rest_data_b2

        # RN choose to suspend the downlink1, DE3->downlink2
        if action == 20 and self.rest_power >= 1:
            # 选择通过downlink2转发给DE3数据包后b3的数据包数量的变化
            self.new_rest_data_b3 = max(
                self.rest_data_b3 - min(self.link_state[11], self.rest_data_b3, self.rest_power), 0)
            self.actual_harvest_power = min(
                self.harvest_power, self.emax - self.rest_power)  # 实际吸收的能量
            self.rest_power = min(
                self.rest_power + self.harvest_power - self.consume_power, self.emax)
            self.reward = self.rw_d * \
                (self.rest_data_b3 - self.new_rest_data_b3)
            self.rest_data_b3 = self.new_rest_data_b3

        # RN choose to suspend the downlink2, DE1->downlink1
        if action == 21 and self.rest_power >= 1:
            # 选择通过downlink1转发给DE1数据包后b1的数据包数量的变化
            self.new_rest_data_b1 = max(
                self.rest_data_b1 - min(self.link_state[6], self.rest_data_b1, self.rest_power), 0)
            self.actual_harvest_power = min(
                self.harvest_power, self.emax - self.rest_power)  # 实际吸收的能量
            self.rest_power = min(
                self.rest_power + self.harvest_power - self.consume_power, self.emax)
            self.reward = self.rw_d * \
                (self.rest_data_b1 - self.new_rest_data_b1)
            self.rest_data_b1 = self.new_rest_data_b1

        # RN choose to suspend the downlink1, DE2->downlink1
        if action == 22 and self.rest_power >= 1:
            # 选择通过downlink1转发给DE2数据包后b2的数据包数量的变化
            self.new_rest_data_b2 = max(
                self.rest_data_b2 - min(self.link_state[8], self.rest_data_b2, self.rest_power), 0)
            self.actual_harvest_power = min(
                self.harvest_power, self.emax - self.rest_power)  # 实际吸收的能量
            self.rest_power = min(
                self.rest_power + self.harvest_power - self.consume_power, self.emax)
            self.reward = self.rw_d * \
                (self.rest_data_b2 - self.new_rest_data_b2)
            self.rest_data_b2 = self.new_rest_data_b2

        # RN choose to suspend the downlink1, DE3->downlink1
        if action == 23 and self.rest_power >= 1:
            # 选择通过downlink1转发给DE3数据包后b3的数据包数量的变化
            self.new_rest_data_b3 = max(
                self.rest_data_b3 - min(self.link_state[10], self.rest_data_b3, self.rest_power), 0)
            self.actual_harvest_power = min(
                self.harvest_power, self.emax - self.rest_power)  # 实际吸收的能量
            self.rest_power = min(
                self.rest_power + self.harvest_power - self.consume_power, self.emax)
            self.reward = self.rw_d * \
                (self.rest_data_b3 - self.new_rest_data_b3)

        # action = wait. RN choose to suspend all the links and wait for more energy
        if action == 24:
            self.reward = 0
            self.actual_harvest_power = min(
                self.harvest_power, self.emax - self.rest_power)  # 实际吸收的能量
            self.rest_power = min(
                self.rest_power + self.actual_harvest_power, self.emax)

        self.link_state = np.zeros(2 * self.num_transmitter * self.num_channel)
        xs = np.random.uniform(0, 1, size=self.link_state.shape)
        for i in range(1, len(self.link_probabilities)):
            self.link_state[(self.link_probabilities_cdf[i-1] < xs) &
                            (xs < self.link_probabilities_cdf[i])] = i

        self.step_counter += 1
        if self.step_counter == self.max_step:
            self.done = True
        else:
            self.done = False

        rest_state = np.array([self.rest_data_b1, self.rest_data_b2,
                               self.rest_data_b3, self.rest_power])
        obs = dict(
            observation=np.concatenate([self.link_state/3, rest_state/10]),
            link_state=self.link_state.astype(np.int),
            rest_datab1=int(self.rest_data_b1),
            rest_datab2=int(self.rest_data_b2),
            rest_datab3=int(self.rest_data_b3),
            rest_power=int(self.rest_power)
        )
        info = {}
        return obs, self.reward, self.done, info


class NetworkSimpleObsWrapper(gym.ObservationWrapper):
    def __init__(self, env: NetworkEnv):
        super().__init__(env)
        self.observation_space = self.env.observation_space['observation']

    def observation(self, observation: Dict[str, np.ndarray]):
        obs = observation['observation']
        return obs
